package icmp

import (
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"go.uber.org/ratelimit"
	"golang.org/x/net/context"
	"net"
	"nprobe/core"
	"nprobe/enviroments"
	"nprobe/utils/collections"
	"nprobe/utils/routines"
	"time"
)

type ICMPSender struct {
	core.BaseSender
}

func NewICMPSender(
	ctx context.Context, name string, id int, dstIPs []string, task *core.ProbeTask, salts map[int][]byte, receiverCnt int,
) *ICMPSender {

	ppRBuf := core.NewPPRBuf(task.SpanMs, task.TimeoutMs)
	copiedSalts := collections.DeepCopyIntMap(salts)

	return &ICMPSender{
		BaseSender: core.BaseSender{
			Ctx:         ctx,
			Name:        name,
			ID:          id,
			DstIPs:      dstIPs,
			Task:        task,
			RBuf:        ppRBuf,
			Salts:       copiedSalts,
			ReceiverCnt: receiverCnt,
		},
	}
}

func (s *ICMPSender) StartSender() error {
	routines.GoWithRecover(s.Name, s.doSend)
	return nil
}

func (s *ICMPSender) doSend() {

	golog.Debugf("run sender. name: %s", s.Name)

	rc, err := core.OpenSenderIP4RawPacketConn("ip4:icmp", "0.0.0.0")
	if err != nil {
		return
	}
	s.Conn = rc

	limiter := ratelimit.New(s.Task.PPS, ratelimit.Per(time.Second))
	var cnt uint64
	var seq uint16
	s.CurrSpanRecord = core.NewSpanRecord(0)
	for {
		for _, srcIP := range s.Task.SrcIPs {
			limiter.Take()

			seq++
			for _, dstIP := range s.DstIPs {
				select {
				case <-s.Ctx.Done():
					s.StopSender()
					return
				default:
					cnt++
					s.sendOnePkt(srcIP, dstIP, seq, cnt)
				}
			}
		}
	}
}

func (s *ICMPSender) sendOnePkt(srcIP string, dstIP string, seq uint16, cnt uint64) error {
	receiverID := cnt % uint64(s.ReceiverCnt)
	payload := s.Salts[int(seq)%4]
	txTs := time.Now().UnixNano()
	binary.BigEndian.PutUint16(payload[8:], uint16(receiverID))
	binary.BigEndian.PutUint16(payload[10:], uint16(0))
	binary.BigEndian.PutUint64(payload[12:], uint64(txTs))

	pid := enviroments.PID
	msg, err := core.EncodeIPICMPEchoRequestPacket(srcIP, dstIP, s.Task.TOS, s.Task.TTL, uint16(pid), seq, payload)
	if err != nil {
		return err
	}
	err = s.Conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	if err != nil {
		return err
	}
	_, err = s.Conn.WriteToIP(msg, &net.IPAddr{IP: net.ParseIP(dstIP)})
	if err != nil {
		return err
	}

	s.recordPkt(srcIP, dstIP, seq, txTs)
	return nil
}

func (s *ICMPSender) recordPkt(srcIP string, dstIP string, seq uint16, txTs int64) {
	spanID := txTs / (s.Task.SpanMs * int64(time.Millisecond))
	golog.Debugf("%s: %s->%s, seq: %d, spanID: %d, txTs: %d", s.Name, srcIP, dstIP, seq, spanID, txTs)
	if spanID > s.CurrSpanRecord.SpanID {
		s.CurrSpanRecord = core.NewSpanRecord(spanID)
		s.RBuf.Enqueue(s.CurrSpanRecord)
	}

	key := fmt.Sprintf("%s-%d", dstIP, seq)
	s.CurrSpanRecord.PktRecordMap[key] = &core.PktRecord{
		Key:   key,
		SrcIP: srcIP,
		DstIP: dstIP,
		Proto: s.Task.Proto,
		Tos:   s.Task.TOS,
		TxTs:  txTs,
	}
}
